import mplhep
import matplotlib.pyplot as plt
plt.style.use(mplhep.style.LHCb2)

import pandas as pd
import numpy as np

def plot_ant_time(file_path, save_path, end_time, param_list):

    fig = plt.figure(figsize=[12,6])

    df = pd.read_csv(file_path)
    time_data = np.array(df[['I','Time']])    

    plt.clf()
    plt.xlim(time_data[0,0], time_data[-1,0])
    plt.ylim(np.min(time_data[:,1])-5, np.max(time_data[:,1])+5)

    ax = fig.get_axes()[0]
    ax.set_xlabel("Iteration")
    ax.set_ylabel("Time / day")

    plt.scatter(time_data[:,0], time_data[:,1], s=5, marker='.', color='black', label='Iteration best ant')

    # 填最优蚂蚁变化数组
    best_time = 1e10
    best_ant_hist = []
    new_best_ant = []
    for itime in time_data:
        if best_time > itime[1]:
            best_time = itime[1] 
            new_best_ant.append([itime[0], best_time])
        best_ant_hist.append([itime[0], best_time])
    best_ant_hist = np.array(best_ant_hist, dtype=np.float32)
    new_best_ant = np.array(new_best_ant, dtype=np.float32)

    plt.plot(best_ant_hist[:,0], best_ant_hist[:,1], 'r-', linewidth=3, label='Global best ant', zorder=-1)
    plt.scatter(new_best_ant[:,0], new_best_ant[:,1], s=50, marker='*', color='Green', label='New Global Best')
    plt.legend()

    plt.savefig(f"{save_path}/acs_ant_time_scatter.png")

    # 计算直方图
    hist, bins = np.histogram(time_data[:,1], 100)
    errors = np.sqrt(hist)
    bin_width = bins[1] - bins[0]
    bin_centers = (bins[:-1] + bins[1:]) / 2

    fig = plt.figure(figsize=[12,6])
    plt.xlim(np.min(time_data[:,1])-5, np.max(time_data[:,1])+5)
    plt.ylim(0, np.max(hist)*1.2)
    # 画直方图
    plt.errorbar(bin_centers, hist, xerr=bin_width/2., yerr=errors, fmt='none', ecolor='black', capsize=1., elinewidth=1.5, label="SOM")
    #plt.legend(fontsize='x-large')

    ax = fig.get_axes()[0]
    ax.set_xlabel("Task time / day")
    ax.set_ylabel("Candidates")

    plt.savefig(f"{save_path}/acs_ant_time_distribution.png")

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Test Operation module")
    #parser.add_argument("-f", "--file", default="result.csv", dest="file")
    parser.add_argument(dest="file", type=str)
    parser.add_argument("-s", "--save_path", default="default", dest="save_path")
    parser.add_argument("-t", "--end_time", type=float, default=-1, dest="end_time")
    arglist = parser.parse_args()
    
    import os
    csv_file = arglist.file
    if not os.path.exists(csv_file):
        raise FileNotFoundError(f"No such file or directory: '{csv_file}'.")
    save_path = csv_file[:csv_file.rfind("/")] + '/ant_time'
    if arglist.save_path != "default":
        save_path = arglist.save_path
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    plot_ant_time(file_path=arglist.file, save_path=save_path, end_time=arglist.end_time, param_list=['ra', 'dec'])

if __name__ == "__main__":
    main()